# 文件路径: BigoImageTaskControl.py
# [最终稳健版] 修正了_upload_image_to_comfyui函数，使其能够智能查找子目录或根目录下的输入文件，以兼容所有工作流。

import uuid
import requests
import json
import time
import os
import traceback
import tempfile
import random
import threading
from datetime import datetime
from pathlib import Path
from fastapi import UploadFile

# 新增的导入，用于同步WebSocket通信
import websocket

import mxupy as mu
import urllib.parse
import bigOAINet as bigo
from mxupy import IM


class BigoImageTaskControl(mu.EntityXControl):
    """
    AI绘画任务控制器 (已修正后台执行逻辑, 新增AI换装功能, 优化服务器选择和数据流)
    """

    class Meta:
        model_class = bigo.ImageGenerateTask

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.comfy_config = mu.read_config().get("comfyui_server", {})
        self.comfy_host = self.comfy_config.get("host")
        self.comfy_port = self.comfy_config.get("port")
        self.comfy_server_url = (
            f"http://{self.comfy_host}:{self.comfy_port}"
            if self.comfy_host and self.comfy_port
            else None
        )
        self.wear_url = (
            "https://d3e6-2001-da8-d00a-fffb-f492-4bad-bf70-faae.ngrok-free.app"
        )
        self.image_control = bigo.BigoImageControl.inst()
        self.api_server_config = mu.read_config().get("api_server", {})
        api_host = self.api_server_config.get("host", "127.0.0.1")
        api_port = self.api_server_config.get("port", 80)
        self.api_base_url = f"http://{api_host}:{api_port}"

    def addBigoImageTask(self, data):
        """
        [业务API] 添加一个新的AI绘画任务。
        此版本修正了数据库事务提交与后台线程启动之间的竞速问题。
        """
        input_image_urls_str = getattr(data, "inputImageUrls", None)
        image_urls_for_thread = []
        if input_image_urls_str:
            try:
                image_urls_for_thread = [
                    url for url in json.loads(input_image_urls_str) if url
                ]
            except Exception as e:
                return IM(False, f"解析输入图片URL时出错: {e}")

        if hasattr(data, "inputImageUrls"):
            delattr(data, "inputImageUrls")

        def _create_task_in_db():
            data.status = "pending"
            im_add = self.add(data)
            if im_add.error:
                return im_add

            task = im_add.data
            for url in image_urls_for_thread:
                self.image_control.add(
                    {
                        "userId": data.userId,
                        "bigoTaskId": task.bigoTaskId,
                        "imageUrl": url,
                        "imageType": "input",
                    }
                )
            return im_add

        im_task_created = self.run(_create_task_in_db)

        if im_task_created.success:
            task = im_task_created.data
            print(
                f"[任务启动] 任务ID: {task.bigoTaskId}, 类型: {task.workflowType}。数据库事务已提交。"
            )
            thread = threading.Thread(
                target=self.generateBigoImage_Batch,
                args=(task.bigoTaskId, data.userId, image_urls_for_thread),
            )
            thread.start()
            print(f"[任务启动] ID {task.bigoTaskId} 的后台线程已派出。")

        return im_task_created

    def generateBigoImage_Batch(self, bigoTaskId, userId, input_image_urls):
        """[内部核心流程] 执行批量AI绘画任务的完整流程。"""

        def _do():
            try:
                print(
                    f"[终极调试] 1. 任务ID {bigoTaskId} 的后台线程已启动，进入_do函数。"
                )
                im_get = self.get_one_by_id(bigoTaskId)
                if im_get.error or not im_get.data:
                    print(f"[终极调试] 错误点：在数据库中找不到任务ID {bigoTaskId}。")
                    return IM(False, f"任务 {bigoTaskId} 不存在。")

                task = im_get.data
                print(
                    f"[终极调试] 2. 成功从数据库获取任务数据，任务类型: {task.workflowType}。"
                )
                self.update_by_id(bigoTaskId, {"status": "processing"})
                print(f"[终极调试] 3. 已将任务状态更新为 'processing'。")

                initial_seed = (
                    int(task.seed)
                    if str(task.seed).isdigit() and task.seed
                    else random.randint(0, 2**32 - 1)
                )
                batch_size = (
                    task.batch_size if task.batch_size and task.batch_size > 0 else 1
                )
                print(
                    f"[终极调试] 4. 参数解析完毕，准备进入循环，总共 {batch_size} 次。"
                )
                for i in range(batch_size):
                    current_seed = initial_seed + i
                    print(
                        f"[终极调试] 5. 准备进入循环，调用 _execute_single_workflow (第{i+1}/{batch_size}次)"
                    )
                    im_exec = self._execute_single_workflow(
                        task, userId, current_seed, input_image_urls
                    )
                    print(f"[终极调试] 6. _execute_single_workflow 已成功返回。")
                    if im_exec.error:
                        print(
                            f"[任务 {bigoTaskId}] 第 {i+1} 张图片生成失败: {im_exec.msg}"
                        )
                        error_details = {
                            "error": f"第 {i+1} 张图片失败: {im_exec.msg}",
                            "traceback": str(im_exec.data),
                        }
                        self.update_by_id(
                            bigoTaskId,
                            {
                                "status": "failed",
                                "finishTime": datetime.now(),
                                "extraParams": json.dumps(error_details),
                            },
                        )
                        return im_exec

                    self.image_control.add(
                        {
                            "userId": userId,
                            "bigoTaskId": bigoTaskId,
                            "imageUrl": im_exec.data,
                            "imageType": "output",
                        }
                    )
                    print(
                        f"[任务 {bigoTaskId}] 第 {i+1} 张输出图片记录已保存: {im_exec.data}"
                    )

                print(f"[任务 {bigoTaskId}] 所有图片生成完毕。")
                self.update_by_id(
                    bigoTaskId, {"status": "completed", "finishTime": datetime.now()}
                )
                return IM(True, "任务全部完成")

            except Exception as e:
                tb_str = traceback.format_exc()
                print("=" * 60)
                print(f"[终极调试] !!! 后台线程在调用工作流前崩溃 !!!")
                print(f"任务ID: {bigoTaskId}")
                print(f"错误类型: {type(e).__name__}")
                print(f"错误信息: {e}")
                print(f"完整追溯信息 (请将这部分完整复制给我):\n{tb_str}")
                print("=" * 60)
                self.update_by_id(
                    bigoTaskId,
                    {
                        "status": "failed",
                        "finishTime": datetime.now(),
                        "extraParams": json.dumps(
                            {"error": str(e), "traceback": tb_str}
                        ),
                    },
                )
                return IM(False, msg=str(e), data=tb_str)

        return self.run(_do)

    def getBigoImageTask(self, bigoTaskId, userId, accessToken):
        def _do():
            im_user_check = bigo.UserControl.inst().check_accesstoken(
                userId, accessToken
            )
            if im_user_check.error:
                return im_user_check
            im_task = self.get_one(
                where=[{"bigoTaskId": bigoTaskId}, {"userId": userId}], to_dict=True
            )
            if im_task.error or not im_task.data:
                return im_task
            im_images = self.image_control.get_list(
                where={"bigoTaskId": bigoTaskId}, to_dict=True
            )
            task_data = im_task.data
            task_data["input_images"] = [
                img for img in im_images.data if img["imageType"] == "input"
            ]
            task_data["output_images"] = [
                img for img in im_images.data if img["imageType"] == "output"
            ]
            im_task.data = task_data
            return im_task

        return self.run(_do)

    def _wait_for_result_ws_sync(
        self, server_url, prompt_id, client_id, save_image_node_id, timeout=300
    ):
        """
        [纯同步版] 使用 websocket-client 库监听 ComfyUI 任务结果。
        此版本无需 asyncio，可安全用于任何同步框架。
        """
        ws_url = server_url.replace("http", "ws", 1)
        ws_uri = f"{ws_url}/ws?clientId={client_id}"
        print(f"[WebSocket-SYNC] 准备连接到: {ws_uri}")

        ws = None
        try:
            # 建立一个同步的、阻塞式的 WebSocket 连接
            ws = websocket.create_connection(ws_uri, timeout=timeout)
            print(
                f"[WebSocket-SYNC] 连接成功 (Prompt ID: {prompt_id})。等待任务完成信号..."
            )

            while True:
                # recv() 是一个阻塞操作，会一直等待直到收到消息
                message_str = ws.recv()
                if not message_str:
                    continue

                message = json.loads(message_str)

                # 监听逻辑和之前完全一样
                if message.get("type") == "executing":
                    data = message.get("data", {})
                    if data.get("node") is None and data.get("prompt_id") == prompt_id:
                        print(f"[WebSocket-SYNC] 收到任务 {prompt_id} 完成信号！")
                        break  # 任务完成，跳出循环

        except websocket.WebSocketTimeoutException:
            raise TimeoutError(f"WebSocket 等待任务 {prompt_id} 结果超时。")
        except Exception as e:
            tb_str = traceback.format_exc()
            print(f"[WebSocket-SYNC] 连接或处理过程中发生严重错误: {e}\n{tb_str}")
            raise e
        finally:
            # 确保无论如何都关闭连接
            if ws:
                ws.close()
                print("[WebSocket-SYNC] 连接已关闭。")

        # --- 后续获取结果的逻辑保持不变 ---
        print(
            f"[WebSocket-SYNC] 信号已收到，正在通过HTTP获取最终结果: /history/{prompt_id}"
        )
        history_url = f"{server_url}/history/{prompt_id}"
        response = requests.get(history_url)
        response.raise_for_status()
        history = response.json()

        if prompt_id in history and history[prompt_id].get("outputs"):
            outputs = history[prompt_id]["outputs"]
            if save_image_node_id in outputs:
                print(
                    f"[WebSocket-SYNC] 成功从History API获取到节点 {save_image_node_id} 的输出！"
                )
                return outputs[save_image_node_id]

        raise ValueError(f"任务 {prompt_id} 在历史记录中未找到或没有输出。")

    def _execute_single_workflow(self, task, userId, seed, input_image_urls):
        try:
            print("[调试] 1. 进入 _execute_single_workflow 方法。")

            target_server_url = self.comfy_server_url
            if (
                task.workflowType == "clothing_swap"
                or task.workflowType == "style_transfer"
            ):
                target_server_url = self.wear_url
                print(
                    f"[调试] 检测到 '{task.workflowType}' 类型，服务器切换至: {target_server_url}"
                )

            elif task.workflowType == "upscale":
                target_server_url = self.comfy_server_url
                print(
                    f"[调试] 检测到 'upscale' 类型，使用主服务器: {target_server_url}"
                )

            print(f"[调试] 2. 最终目标服务器地址: {target_server_url}")
            if not target_server_url:
                raise ValueError("目标ComfyUI服务器地址未在配置文件中正确设置。")

            project_root = Path(__file__).resolve().parents[2]
            workflow_filename = f"{task.workflowType}.json"
            workflow_path = os.path.join(
                project_root, "static", "workflows", workflow_filename
            )

            # 兼容 clothing_swap 可能使用旧文件名的情况
            if (
                not os.path.exists(workflow_path)
                and task.workflowType == "clothing_swap"
            ):
                workflow_path = os.path.join(
                    project_root, "static", "workflows", "Replacement.json"
                )

            print(f"[调试] 3. 正在尝试加载工作流文件: {workflow_path}")
            if not os.path.exists(workflow_path):
                raise FileNotFoundError(f"工作流文件不存在: {workflow_path}")

            with open(workflow_path, "r", encoding="utf-8") as f:
                workflow = json.load(f)
            print(f"[调试] 4. 成功加载工作流文件。")

            node_ids = self._find_node_ids(workflow, task.workflowType)
            print(f"[调试] 5. 查找到的关键节点ID: {node_ids}")
            if not node_ids.get("save_image"):
                raise ValueError(
                    "关键的 'SaveImage' 节点未在工作流中找到，请检查JSON文件。"
                )

            print("[调试] 6. 准备向ComfyUI发送数据...")
            self._prepare_workflow_inputs(
                target_server_url,
                task,
                workflow,
                node_ids,
                userId,
                seed,
                input_image_urls,
            )
            print("[调试] 7. 数据准备完毕，即将发送请求。")

            client_id = str(uuid.uuid4())
            prompt_payload = {"prompt": workflow, "client_id": client_id}

            print(
                f"[调试] 8. 正在向 {target_server_url}/prompt 发送POST请求 (Client ID: {client_id})..."
            )
            response = requests.post(
                f"{target_server_url}/prompt", json=prompt_payload, timeout=10
            )
            print(f"[调试] 9. 收到ComfyUI响应，状态码: {response.status_code}")
            response.raise_for_status()

            prompt_id = response.json()["prompt_id"]
            print(f"[调试] 10. ComfyUI任务已提交，Prompt ID: {prompt_id}")

            print(f"[改造] 调用纯同步方法 _wait_for_result_ws_sync")
            result_data = self._wait_for_result_ws_sync(
                target_server_url, prompt_id, client_id, node_ids.get("save_image")
            )

            final_filename = self._download_and_save_single_result(
                target_server_url, result_data, task
            )

            return IM(True, data=final_filename)
        except Exception as e:
            tb_str = traceback.format_exc()
            print("=" * 60)
            print(f"[调试] !!! 后台线程发生严重错误 !!!")
            print(f"任务ID: {task.bigoTaskId if task else 'N/A'}")
            print(f"错误类型: {type(e).__name__}")
            print(f"错误信息: {e}")
            print(f"完整追溯信息:\n{tb_str}")
            print("=" * 60)
            if task:
                self.update_by_id(
                    task.bigoTaskId,
                    {
                        "status": "failed",
                        "extraParams": json.dumps(
                            {"error": str(e), "traceback": tb_str}
                        ),
                    },
                )
            return IM(False, msg=str(e), data=tb_str)

    def _prepare_workflow_inputs(
        self, server_url, task, workflow, node_ids, userId, seed, input_image_urls
    ):
        print(f"[准备工作流] 种子: {seed}, 类型: {task.workflowType}")
        positive_prompt_text = task.positivePrompt
        negative_prompt_text = task.negativePrompt
        model_name = task.modelName
        extra_params = json.loads(task.extraParams) if task.extraParams else {}

        if task.workflowType == "clothing_swap":
            print("[准备工作流] 检测到 'clothing_swap' 类型，开始进行精确映射...")
            upper_body_ref, lower_body_ref, upper_body_mask, lower_body_mask = (
                input_image_urls
            )
            print("[上传] 开始上传换装所需的图片...")
            img_upper_ref = (
                self._upload_image_to_comfyui(server_url, upper_body_ref, userId, task)
                if upper_body_ref
                else None
            )
            img_lower_ref = (
                self._upload_image_to_comfyui(server_url, lower_body_ref, userId, task)
                if lower_body_ref
                else None
            )
            img_upper_mask = (
                self._upload_image_to_comfyui(server_url, upper_body_mask, userId, task)
                if upper_body_mask
                else None
            )
            img_lower_mask = (
                self._upload_image_to_comfyui(server_url, lower_body_mask, userId, task)
                if lower_body_mask
                else None
            )

            if not img_upper_ref or not img_lower_ref or not img_upper_mask:
                raise ValueError("上半身参考图、下半身参考图和上半身蒙版是必需的。")

            workflow["4"]["inputs"]["ckpt_name"] = model_name
            workflow["6"]["inputs"]["text"] = positive_prompt_text
            workflow["7"]["inputs"]["text"] = negative_prompt_text
            workflow["12"]["inputs"]["image"] = img_upper_ref.data["name"]
            workflow["36"]["inputs"]["image"] = img_upper_mask.data["name"]
            workflow["17"]["inputs"]["image"] = img_lower_ref.data["name"]
            workflow["15"]["inputs"]["image"] = (
                img_lower_mask.data["name"]
                if img_lower_mask
                else img_upper_mask.data["name"]
            )
            print("[工作流修改] 已将图片映射到IPAdapter节点。")

            sampler_node = workflow.get("3", {})
            if sampler_node:
                sampler_node["inputs"]["seed"] = seed
                sampler_node["inputs"]["steps"] = int(extra_params.get("steps", 8))
                sampler_node["inputs"]["cfg"] = float(extra_params.get("cfg", 2.0))

            ip_adapter_weight = float(extra_params.get("ip_adapter_weight", 0.8))
            workflow["13"]["inputs"]["weight"] = ip_adapter_weight
            workflow["16"]["inputs"]["weight"] = ip_adapter_weight
            print(
                f"[参数应用] 已应用高级参数: steps={sampler_node['inputs']['steps']}, cfg={sampler_node['inputs']['cfg']}, ip_weight={ip_adapter_weight}"
            )
            return

        elif task.workflowType == "upscale":
            print(
                f"[准备工作流] 检测到 'upscale' 类型，严格遵循 `超分.json` 的预设模型。"
            )

            node_map = {
                "load_image": "2",
                "supir_conditioner": "9",
                "ultimate_sd_upscale": "25",
                "usdu_positive_prompt": "27",
                "seed": "29",
                "usdu_negative_prompt": "31",
            }

            if not input_image_urls:
                raise ValueError("超分任务 (upscale) 缺少输入图片。")
            im_upload = self._upload_image_to_comfyui(
                server_url, input_image_urls[0], userId, task
            )
            if im_upload.error:
                raise Exception(f"上传用于超分的图片到ComfyUI失败: {im_upload.msg}")

            workflow[node_map["load_image"]]["inputs"]["image"] = im_upload.data["name"]
            workflow[node_map["supir_conditioner"]]["inputs"][
                "positive_prompt"
            ] = positive_prompt_text
            workflow[node_map["supir_conditioner"]]["inputs"][
                "negative_prompt"
            ] = negative_prompt_text
            workflow[node_map["usdu_positive_prompt"]]["inputs"][
                "text"
            ] = positive_prompt_text
            workflow[node_map["usdu_negative_prompt"]]["inputs"][
                "text"
            ] = negative_prompt_text
            workflow[node_map["seed"]]["inputs"]["seed"] = seed

            usdu_inputs = workflow[node_map["ultimate_sd_upscale"]]["inputs"]
            if "denoise" in extra_params:
                usdu_inputs["denoise"] = float(extra_params["denoise"])
            if "upscale_by" in extra_params:
                usdu_inputs["upscale_by"] = float(extra_params["upscale_by"])
            if "steps" in extra_params:
                usdu_inputs["steps"] = int(extra_params["steps"])
            if "cfg" in extra_params:
                usdu_inputs["cfg"] = float(extra_params["cfg"])

            print("[准备工作流] `upscale` 类型动态参数注入完成。")
            return

        elif task.workflowType == "style_transfer":
            print("[准备工作流] 检测到 'style_transfer' 类型，开始进行精确映射...")
            if not input_image_urls or len(input_image_urls) < 2:
                raise ValueError(
                    "风格迁移任务 (style_transfer) 需要2个输入图片: [风格图, 内容图]。"
                )

            style_image_url, content_image_url = (
                input_image_urls[0],
                input_image_urls[1],
            )
            print(
                f"[准备工作流] 风格图: {style_image_url}, 内容图: {content_image_url}"
            )

            im_upload_style = self._upload_image_to_comfyui(
                server_url, style_image_url, userId, task
            )
            if im_upload_style.error:
                raise Exception(f"上传风格图片到ComfyUI失败: {im_upload_style.msg}")

            im_upload_content = self._upload_image_to_comfyui(
                server_url, content_image_url, userId, task
            )
            if im_upload_content.error:
                raise Exception(f"上传内容图片到ComfyUI失败: {im_upload_content.msg}")

            node_map = {
                "checkpoint_loader": "4",
                "positive_prompt": "6",
                "negative_prompt": "7",
                "k_sampler": "3",
                "style_image_loader": "12",
                "content_image_loader": "15",
                "ipadapter_style_encoder": "16",
                "ipadapter_content_encoder": "17",
                "ipadapter_main_control": "14",
                "save_image": "9",
            }

            print("[准备工作流] 开始向 'style_transfer' 工作流注入参数...")
            workflow[node_map["checkpoint_loader"]]["inputs"]["ckpt_name"] = model_name
            workflow[node_map["positive_prompt"]]["inputs"][
                "text"
            ] = positive_prompt_text
            workflow[node_map["negative_prompt"]]["inputs"][
                "text"
            ] = negative_prompt_text
            print(f"  -> 注入 [模型: {model_name}] 和 [提示词]")

            workflow[node_map["style_image_loader"]]["inputs"]["image"] = (
                im_upload_style.data["name"]
            )
            workflow[node_map["content_image_loader"]]["inputs"]["image"] = (
                im_upload_content.data["name"]
            )
            print(
                f"  -> 注入 [风格图: {im_upload_style.data['name']}] 和 [内容图: {im_upload_content.data['name']}]"
            )

            sampler_node = workflow[node_map["k_sampler"]]
            sampler_node["inputs"]["seed"] = seed
            sampler_node["inputs"]["steps"] = int(extra_params.get("steps", 30))
            sampler_node["inputs"]["cfg"] = float(extra_params.get("cfg", 6.8))
            print(
                f"  -> 注入 [采样器参数] Seed: {seed}, Steps: {sampler_node['inputs']['steps']}, CFG: {sampler_node['inputs']['cfg']}"
            )

            main_control_node = workflow[node_map["ipadapter_main_control"]]
            main_control_node["inputs"]["weight"] = float(
                extra_params.get("main_weight", 0.8)
            )
            main_control_node["inputs"]["start_at"] = float(
                extra_params.get("start_at", 0.0)
            )
            main_control_node["inputs"]["end_at"] = float(
                extra_params.get("end_at", 1.0)
            )

            workflow[node_map["ipadapter_style_encoder"]]["inputs"]["weight"] = float(
                extra_params.get("style_weight", 0.75)
            )
            workflow[node_map["ipadapter_content_encoder"]]["inputs"]["weight"] = float(
                extra_params.get("content_weight", 1.0)
            )

            print(
                f"  -> 注入 [IPAdapter权重] Style: {extra_params.get('style_weight', 0.75)}, Content: {extra_params.get('content_weight', 1.0)}, Main: {extra_params.get('main_weight', 0.8)}"
            )
            print(
                f"  -> 注入 [IPAdapter步数] Start: {extra_params.get('start_at', 0.0)}, End: {extra_params.get('end_at', 1.0)}"
            )

            print("[准备工作流] 'style_transfer' 工作流参数注入完成。")
            return

        workflow[node_ids["checkpoint_loader"]]["inputs"]["ckpt_name"] = model_name
        if node_ids.get("negative_prompt"):
            workflow[node_ids["negative_prompt"]]["inputs"][
                "text"
            ] = negative_prompt_text
        if node_ids.get("k_sampler"):
            workflow[node_ids["k_sampler"]]["inputs"]["seed"] = seed

        if task.workflowType in [
            "text_to_image",
            "lora",
            "lora_multiple",
            "embedding",
            "gligen",
        ]:
            if node_ids.get("empty_latent"):
                latent_inputs = workflow[node_ids["empty_latent"]]["inputs"]
                latent_inputs["width"] = task.width
                latent_inputs["height"] = task.height
                latent_inputs["batch_size"] = 1

        elif task.workflowType in ["image_to_image", "inpaint", "outpaint"]:
            if not node_ids.get("load_image"):
                raise ValueError(f"{task.workflowType} 工作流缺少'LoadImage'节点")
            if not input_image_urls:
                raise ValueError(f"{task.workflowType} 任务缺少输入图片")
            im_upload = self._upload_image_to_comfyui(
                server_url, input_image_urls[0], userId, task
            )
            if im_upload.error:
                raise Exception(f"上传输入图片到ComfyUI失败: {im_upload.msg}")
            workflow[node_ids["load_image"]]["inputs"]["image"] = im_upload.data["name"]

        if task.workflowType == "lora":
            if not node_ids.get("lora_loader"):
                raise ValueError("LoRA工作流缺少'LoraLoader'节点")
            lora_node_inputs = workflow[node_ids["lora_loader"]]["inputs"]
            lora_node_inputs["lora_name"] = extra_params.get("lora_name")
            if not lora_node_inputs["lora_name"]:
                raise ValueError("使用LoRA工作流时，extraParams必须包含 lora_name")
            lora_node_inputs["strength_model"] = float(
                extra_params.get("strength_model", 0.75)
            )
            lora_node_inputs["strength_clip"] = float(
                extra_params.get("strength_clip", 1.0)
            )

        if node_ids.get("positive_prompt"):
            workflow[node_ids["positive_prompt"]]["inputs"][
                "text"
            ] = positive_prompt_text
            print(f"[任务 {task.bigoTaskId}] 最终正向提示词: {positive_prompt_text}")

    def _find_node_ids(self, workflow, workflow_type):
        node_ids = {}
        if workflow_type == "style_transfer":
            node_ids["save_image"] = "9"
            return node_ids
        if workflow_type == "clothing_swap":
            node_ids["save_image"] = "32"
            return node_ids
        if workflow_type == "multi-image_fusion":
            node_ids["save_image"] = "67"
        for nid, node in workflow.items():
            class_type = node.get("class_type")
            if not class_type:
                continue

            key_map = {
                "CheckpointLoaderSimple": "checkpoint_loader",
                "KSampler": "k_sampler",
                "SaveImage": "save_image",
                "VAEDecode": "vae_decode",
                "EmptyLatentImage": "empty_latent",
                "LoadImage": "load_image",
                "VAEEncodeForInpaint": "vae_encode_inpaint",
                "ImagePadForOutpaint": "image_pad_for_outpaint",
                "GLIGENLoader": "gligen_loader",
                "LoraLoader": "lora_loader",
                "UltimateSDUpscale": "ultimate_sd_upscale",
                "UpscaleModelLoader": "upscale_model_loader",
            }

            if class_type in key_map:
                node_ids[key_map[class_type]] = nid

        sampler_node_id = node_ids.get("k_sampler")
        if sampler_node_id and workflow.get(sampler_node_id):
            sampler_inputs = workflow[sampler_node_id].get("inputs", {})
            if "positive" in sampler_inputs:
                node_ids["positive_prompt"] = sampler_inputs["positive"][0]
            if "negative" in sampler_inputs:
                node_ids["negative_prompt"] = sampler_inputs["negative"][0]

        upscaler_node_id = node_ids.get("ultimate_sd_upscale")
        if upscaler_node_id and workflow.get(upscaler_node_id):
            upscaler_inputs = workflow[upscaler_node_id].get("inputs", {})
            if "positive" in upscaler_inputs and "positive_prompt" not in node_ids:
                node_ids["positive_prompt"] = upscaler_inputs["positive"][0]
            if "negative" in upscaler_inputs and "negative_prompt" not in node_ids:
                node_ids["negative_prompt"] = upscaler_inputs["negative"][0]

        return node_ids

    # [最终修正] 重构此函数以同时兼容两种文件路径
    def _upload_image_to_comfyui(
        self, server_url, relative_image_path, userId, task=None
    ):
        user_file_root = mu.read_config().get("api_server", {}).get("user_file_path")
        if not user_file_root:
            raise ValueError("配置文件中缺少 user_file_path")

        # [稳健性重构] 设计一个能适应两种情况的健壮的文件查找逻辑
        absolute_image_path = None

        # 步骤1: 优先尝试在工作流对应的子目录中查找
        # 兼容 clothing_swap 这类将输入文件预置在子目录的情况
        # 也兼容某些前端页面上传时就指定了子目录的情况
        possible_sub_dirs = []
        if task and task.workflowType:
            possible_sub_dirs.append(task.workflowType)
            if task.workflowType == "clothing_swap":
                possible_sub_dirs.append("swap")  # 兼容旧的 "swap" 目录

        for sub_dir in possible_sub_dirs:
            path_in_subdir = os.path.join(
                user_file_root, str(userId), sub_dir, relative_image_path
            )
            print(f"[文件查找] 步骤1: 尝试在子目录 '{sub_dir}' 查找: {path_in_subdir}")
            if os.path.exists(path_in_subdir):
                absolute_image_path = path_in_subdir
                print(f"[文件查找] 成功在子目录 '{sub_dir}' 找到文件。")
                break

        # 步骤2: 如果在所有可能的子目录中都没找到，则回退到在用户根目录中查找
        if absolute_image_path is None:
            path_in_root = os.path.join(
                user_file_root, str(userId), relative_image_path
            )
            print(f"[文件查找] 步骤2: 尝试在根目录查找: {path_in_root}")
            if os.path.exists(path_in_root):
                absolute_image_path = path_in_root
                print(f"[文件查找] 成功在根目录找到文件。")

        # 步骤3: 如果两种方式都找不到，才最终报错
        if absolute_image_path is None:
            final_attempted_path = os.path.join(
                user_file_root, str(userId), relative_image_path
            )
            raise FileNotFoundError(
                f"输入文件在所有可预见的路径均未找到。根目录路径: {final_attempted_path}"
            )

        # --- 后续上传逻辑不变 ---
        print(f"[上传] 准备上传文件: {absolute_image_path} 到 {server_url}")
        with open(absolute_image_path, "rb") as f:
            files = {"image": (os.path.basename(relative_image_path), f)}
            response = requests.post(
                f"{server_url}/upload/image", files=files, data={"overwrite": "true"}
            )
            response.raise_for_status()
            print(f"[上传] 文件 {os.path.basename(relative_image_path)} 上传成功。")
            return IM(True, data=response.json())

    def _download_and_save_single_result(self, server_url, result_data, task):
        output_images = result_data.get("images", [])
        if not output_images:
            raise ValueError("ComfyUI任务完成，但未返回图片。")
        image_info = output_images[0]
        filename = image_info.get("filename")
        if not filename:
            raise ValueError("ComfyUI返回的结果中缺少'filename'。")
        subfolder = image_info.get("subfolder", "")
        img_type = image_info.get("type", "output")
        view_url = f"{server_url}/view?filename={urllib.parse.quote(filename)}&subfolder={urllib.parse.quote(subfolder)}&type={img_type}"
        response = requests.get(view_url, stream=True)
        response.raise_for_status()
        with tempfile.NamedTemporaryFile(
            delete=True, suffix=f"_{filename}"
        ) as tmp_file:
            for chunk in response.iter_content(chunk_size=8192):
                tmp_file.write(chunk)
            tmp_file.seek(0)
            fastapi_upload_file = UploadFile(file=tmp_file, filename=filename)
            # [正确逻辑] 此处为保存输出图片，使用子目录是正确的
            sub_directory = task.workflowType
            im_upload = mu.upload_user_file(
                file=fastapi_upload_file,
                user_id=task.user.userId,
                sub_dir=sub_directory,
                access_token="A_VALID_INTERNAL_TOKEN_IF_NEEDED",
                keep=True,
                override=True,
            )
        if im_upload.error:
            raise Exception(f"调用框架文件服务保存结果图片失败: {im_upload.msg}")
        new_filename = (
            im_upload.data.get("filename")
            if isinstance(im_upload.data, dict)
            else str(im_upload.data)
        )
        if not new_filename:
            raise Exception("框架文件服务未返回有效文件名。")
        return new_filename

    def getHistory(self, data):
        def _do():
            userId = getattr(data, "userId", None)
            accessToken = getattr(data, "accessToken", None)
            limit = getattr(data, "limit", 20)
            offset = getattr(data, "offset", 0)
            im_user_check = bigo.UserControl.inst().check_accesstoken(
                userId, accessToken
            )
            if im_user_check.error:
                return im_user_check
            im_tasks = self.get_list(
                where={"userId": userId},
                order_by=[{"addTime": "desc"}],
                limit=limit,
                offset=offset,
                to_dict=True,
            )
            if im_tasks.error:
                return im_tasks
            tasks_list = im_tasks.data
            task_ids = [task["bigoTaskId"] for task in tasks_list]
            if not task_ids:
                return IM(True, "历史记录为空", data=[])
            im_images = self.image_control.get_list(
                where={"bigoTaskId__in": task_ids, "imageType": "output"},
                to_dict=True,
                recurse=True,
            )
            images_data = im_images.data if im_images.success else []
            images_by_task_id = {}
            for image in images_data:
                task_id = image["task"]["bigoTaskId"]
                workflow_type = image["task"]["workflowType"]
                filename = image["imageUrl"]
                image["imageUrl"] = (
                    f"{self.api_base_url}/file?filename={urllib.parse.quote(filename)}&sub_dir={workflow_type}&userId={userId}&access_token={accessToken}"
                )
                if task_id not in images_by_task_id:
                    images_by_task_id[task_id] = []
                images_by_task_id[task_id].append(image)
            for task in tasks_list:
                task["output_images"] = images_by_task_id.get(task["bigoTaskId"], [])
            im_tasks.data = tasks_list
            return im_tasks

        return self.run(_do)
